#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
pick_fastest_conda.py
自动测速常见 conda 通道（defaults 与镜像），选择最快组合，并可持久化写入用户级 ~/.condarc。

设计要点：
- 仅用标准库（urllib、ssl、time、subprocess、argparse、platform），零外部依赖。
- 以真实路径测速度：访问 <base>/pkgs/<main|r>/<subdir>/repodata.json（以及可选 conda-forge）。
- 指标：单通道用“连接 + 首包 + 连续读取上限 bytes”的总耗时；多次尝试取最好；多通道求和作为组合评分。
- 持久化：使用 `conda config` 修改“channels”列表（用户级），不改系统级 /etc/conda/condarc。
  如需回滚到官方默认（defaults）：运行 conda config --remove-key channels 即可。
"""

import argparse
import sys
import time
import urllib.request
import ssl
import subprocess
import platform
from typing import List, Tuple, Dict, Optional

UA = "conda-source-bench/1.0 (+https://docs.conda.io)"
DEFAULT_BYTES = 256 * 1024
DEFAULT_TRIES = 2
DEFAULT_TIMEOUT = 10.0

# 基础候选（base 即 anaconda 根目录；defaults 的 main/r 均在 base/pkgs/ 之下）
DEFAULT_BASES = [
    # 官方 defaults（最权威，全球可用）
    "https://repo.anaconda.com/anaconda",
    # 常见镜像（可按需增删）
    "https://mirrors.tuna.tsinghua.edu.cn/anaconda",
    "https://mirrors.ustc.edu.cn/anaconda",
    "https://mirrors.aliyun.com/anaconda",
    "https://mirrors.bfsu.edu.cn/anaconda",
]

# conda-forge 的镜像位置（部分镜像站支持 cloud 目录）
FORGE_BY_BASE = {
    "https://repo.anaconda.com/anaconda": "https://conda.anaconda.org/conda-forge",
    "https://mirrors.tuna.tsinghua.edu.cn/anaconda": "https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge",
    "https://mirrors.ustc.edu.cn/anaconda": "https://mirrors.ustc.edu.cn/anaconda/cloud/conda-forge",
    "https://mirrors.aliyun.com/anaconda": None,  # 未稳定提供 conda-forge，留空自动跳过
    "https://mirrors.bfsu.edu.cn/anaconda": "https://mirrors.bfsu.edu.cn/anaconda/cloud/conda-forge",
}

def detect_subdir() -> str:
    """
    推断 conda 子目录（subdir）。常见值：
    linux-64, win-64, osx-64 (Intel), osx-arm64 (Apple Silicon).
    """
    sysplat = sys.platform
    mach = platform.machine().lower()
    if sysplat.startswith("linux"):
        return "linux-64"
    if sysplat.startswith("win"):
        return "win-64"
    if sysplat == "darwin":
        if "arm" in mach or "aarch64" in mach:
            return "osx-arm64"
        return "osx-64"
    # 回退到通用档
    return "linux-64"

def format_secs(s: float) -> str:
    return f"{s*1000:.0f} ms" if s < 1 else f"{s:.2f} s"

def http_read_time(url: str, max_bytes: int, timeout: float) -> Tuple[Optional[float], str]:
    """
    读取最多 max_bytes 字节的耗时（秒）。失败返回 (None, 注记)。
    """
    ctx = ssl.create_default_context()
    start = time.perf_counter()
    first_byte_at = None
    read_bytes = 0
    try:
        req = urllib.request.Request(url, headers={"User-Agent": UA})
        with urllib.request.urlopen(req, timeout=timeout, context=ctx) as resp:
            chunk = resp.read(1)
            if not chunk:
                return None, "empty response"
            first_byte_at = time.perf_counter()
            read_bytes += 1
            while read_bytes < max_bytes:
                c = resp.read(min(8192, max_bytes - read_bytes))
                if not c:
                    break
                read_bytes += len(c)
        elapsed = time.perf_counter() - start
        extra = ""
        if first_byte_at is not None:
            rtt_like = first_byte_at - start
            extra = f" | RTT≈{format_secs(rtt_like)} | read={read_bytes} B"
        return elapsed, f"ok{extra}"
    except Exception as e:
        return None, f"error: {type(e).__name__}: {e}"

def bench_one_base(base: str, subdir: str, max_bytes: int, timeout: float, tries: int,
                   include_forge: bool) -> Tuple[str, Optional[float], List[str], str]:
    """
    基于一个 base（如 https://mirrors.tuna.tsinghua.edu.cn/anaconda ）评估组合耗时：
    - defaults: pkgs/main, pkgs/r
    - 可选：conda-forge 映射（若 FORGE_BY_BASE 有对应项）
    评分采用：各通道（main/r[/forge]）的最好耗时求和。
    返回：(base, 总分 或 None, 实际通道列表, 附加说明)
    """
    channels = [
        f"{base.rstrip('/')}/pkgs/main",
        f"{base.rstrip('/')}/pkgs/r",
    ]
    note_parts = []
    if include_forge:
        forge = FORGE_BY_BASE.get(base)
        if forge:
            channels.append(forge.rstrip('/'))
        else:
            note_parts.append("forge: N/A")

    total = 0.0
    tested = []
    for ch in channels:
        url = f"{ch}/{subdir}/repodata.json"
        best = None
        last_note = ""
        for _ in range(tries):
            t, n = http_read_time(url, max_bytes, timeout)
            last_note = n
            if t is not None:
                best = t if best is None else min(best, t)
        if best is None:
            # 某通道不可用则判整个 base 失败（确保安装不踩坑）
            return base, None, channels, f"fail at {ch} -> {last_note}"
        total += best
        tested.append(ch)
    return base, total, tested, "; ".join(note_parts) if note_parts else "ok"

def rank_bases(bases: List[str], subdir: str, max_bytes: int, timeout: float, tries: int,
               include_forge: bool):
    results = []
    seen = set()
    uniq = []
    for b in bases:
        b = b.strip().rstrip("/")
        if b and b not in seen:
            uniq.append(b)
            seen.add(b)
    for b in uniq:
        base, score, chs, note = bench_one_base(b, subdir, max_bytes, timeout, tries, include_forge)
        results.append((base, score, chs, note))
    results.sort(key=lambda x: float("inf") if x[1] is None else x[1])
    return results

def apply_channels(channels: List[str]) -> Tuple[bool, str]:
    """
    通过 `conda config` 写入用户级 channels（覆盖式：先清空再添加）。
    返回 (成功与否, 输出/错误信息)。
    """
    # 说明（防缩写误解）：`--remove-key` 表示删除配置键；`--add` 表示向列表添加一个值
    cmds = [
        (["conda", "config", "--remove-key", "channels"], False),
    ] + [
        (["conda", "config", "--add", "channels", ch], False) for ch in channels
    ] + [
        (["conda", "config", "--set", "show_channel_urls", "true"], False)
    ]
    out_all = []
    for cmd, _ in cmds:
        try:
            p = subprocess.run(cmd, capture_output=True, text=True, check=False)
        except FileNotFoundError:
            # 某些环境 conda 不是可执行文件（如 Windows 上的 “conda.bat” 或 shell 函数）
            # 尝试以 Python 模块形式调用（注意：部分发行版不支持）
            alt = [sys.executable, "-m", "conda"] + cmd[1:]
            try:
                p = subprocess.run(alt, capture_output=True, text=True, check=False)
            except Exception as e:
                return False, f"spawn error: {type(e).__name__}: {e}"
        out_all.append((cmd, p.returncode, (p.stdout or "") + (p.stderr or "")))
        if cmd[2] == "--remove-key" and out_all[-1][1] != 0:
            # 清空失败可能是因为本就不存在，忽略继续
            pass
        elif p.returncode != 0:
            return False, f"cmd failed: {' '.join(cmd)}\n{out_all[-1][2]}"
    return True, "\n".join(["$ " + " ".join(c) + "\n" + msg for c, rc, msg in out_all])

def show_current_config() -> str:
    try:
        p = subprocess.run(["conda", "config", "--show"], capture_output=True, text=True, check=False)
        if p.returncode != 0 and p.stderr:
            return p.stderr
        return (p.stdout or "") + (p.stderr or "")
    except Exception as e:
        return f"cannot read conda config: {e}"

def main():
    ap = argparse.ArgumentParser(
        description="自动测速并设置最快的 conda 通道（defaults 与镜像）。"
    )
    ap.add_argument("--candidates", type=str, default="",
                    help="自定义候选基础地址（逗号分隔），如：https://repo.anaconda.com/anaconda,https://mirrors.tuna.tsinghua.edu.cn/anaconda")
    ap.add_argument("--include-conda-forge", action="store_true",
                    help="同时测试并配置 conda-forge（若该镜像提供）。")
    ap.add_argument("--subdir", type=str, default="",
                    help="指定平台子目录（默认自动判断），如 linux-64 / win-64 / osx-64 / osx-arm64。")
    ap.add_argument("--bytes", type=int, default=DEFAULT_BYTES,
                    help=f"单通道最大读取字节上限（默认：{DEFAULT_BYTES} 字节）。")
    ap.add_argument("--timeout", type=float, default=DEFAULT_TIMEOUT,
                    help=f"单次连接+读的超时（秒，默认：{DEFAULT_TIMEOUT}）。")
    ap.add_argument("--tries", type=int, default=DEFAULT_TRIES,
                    help=f"每个通道的尝试次数（默认：{DEFAULT_TRIES}，取最好的一次）。")
    ap.add_argument("--apply", action="store_true",
                    help="将最快组合持久化为当前用户默认 channels（使用 conda config）。")
    args = ap.parse_args()

    bases = list(DEFAULT_BASES)
    if args.candidates.strip():
        bases += [x.strip() for x in args.candidates.split(",") if x.strip()]

    subdir = args.subdir.strip() or detect_subdir()

    print("== conda 源测速开始 ==")
    print(f"候选基础地址：{len(bases)} | 子目录：{subdir} | 读取上限：{args.bytes} B | 重试：{args.tries} | forge：{'on' if args.include_conda_forge else 'off'}")

    results = rank_bases(bases, subdir, args.bytes, args.timeout, args.tries, args.include_conda_forge)

    print("\n== 组合速度排名（越上越快）==")
    best = None
    best_channels = None
    for i, (base, score, chs, note) in enumerate(results, 1):
        speed = "不可达/失败" if score is None else format_secs(score)
        print(f"{i:02d}. {base:<45} -> {speed}   [{note}]")
        for ch in chs:
            print(f"     - {ch}")
        if best is None and score is not None:
            best = base
            best_channels = chs

    if best is None:
        print("\n未找到可用组合。请检查网络、代理（proxy）或更换候选列表。")
        sys.exit(2)

    print(f"\n最快基础：{best}")
    print("通道组成（写入顺序同下）：")
    for ch in best_channels:
        print("  -", ch)

    print("\n== 当前 conda 配置（节选）==")
    print(show_current_config())

    if args.apply:
        print("\n== 写入默认 channels ==")
        ok, out = apply_channels(best_channels)
        if ok:
            print("已将默认 channels 更新为：")
            for ch in best_channels:
                print("  -", ch)
        else:
            print("写入失败：")
        print(out)

        print("\n== 更新后 conda 配置（节选）==")
        print(show_current_config())
    else:
        print("\n（仅测速未写入。如需持久化，追加参数：--apply）")

if __name__ == "__main__":
    main()
